import paddlex as pdx

insect_dataset = 'https://bj.bcebos.com/paddlex/datasets/insect_det.tar.gz'
pdx.utils.download_and_decompress(insect_dataset, path='./')


from paddlex import transforms as T

train_transforms = T.Compose([
    T.MixupImage(mixup_epoch=250), T.RandomDistort(),
    T.RandomExpand(im_padding_value=[123.675, 116.28, 103.53]), T.RandomCrop(),
    T.RandomHorizontalFlip(), T.BatchRandomResize(
        target_sizes=[320, 352, 384, 416, 448, 480, 512, 544, 576, 608],
        interp='RANDOM'), T.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

eval_transforms = T.Compose([
    T.Resize(
        608, interp='CUBIC'), T.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


# ### 3.2 定义数据集Dataset
# 目标检测可使用`VOCDetection`格式和`COCODetection`两种数据集，此处由于数据集为VOC格式，因此采用`pdx.datasets.VOCDetection`来加载数据集，该接口的介绍可参见文档[paddlex.datasets.VOCDetection](https://github.com/PaddlePaddle/PaddleX/blob/develop/docs/apis/datasets.md#paddlexdatasetsvocdetection)。

# In[ ]:


train_dataset = pdx.datasets.VOCDetection(
    data_dir='insect_det',
    file_list='insect_det/train_list.txt',
    label_list='insect_det/labels.txt',
    transforms=train_transforms,
    shuffle=True)
eval_dataset = pdx.datasets.VOCDetection(
    data_dir='insect_det',
    file_list='insect_det/val_list.txt',
    label_list='insect_det/labels.txt',
    transforms=eval_transforms)


# ### 3.3 模型开始训练
# 使用本数据集在V100上训练，如有GPU，模型的训练过程预估为90分钟左右；如无GPU，则预估为14小时左右。更多训练模型的参数可参见文档[paddlex.det.YOLOv3](https://github.com/PaddlePaddle/PaddleX/blob/develop/docs/apis/models/detection.md#paddlexdetyolov3)。模型训练过程每间隔save_interval_epochs轮会保存一次模型在save_dir目录下，同时在保存的过程中也会在验证数据集上计算相关指标，具体相关日志参见[文档](https://github.com/PaddlePaddle/PaddleX/blob/develop/docs/appendix/metrics.md#%E8%AF%84%E4%BC%B0%E6%97%A5%E5%BF%97%E5%AD%97%E6%AE%B5-1)。  
# 
# **AIStudio使用VisualDL查看训练过程中的指标变化**  
# 1. 点击左边菜单图标的『**可视化**』；  
# 2. 设置logdir，logdir的路径为训练代码中`save_dir`指定的目录下的vdl_log目录，例如`output/yolov3_darknet53/vdl_log`  
# 3. 点击下方『**启动VisualDL服务按钮**』，再『**打开VisualDL**』即可

# In[3]:


num_classes = len(train_dataset.labels)
model = pdx.det.YOLOv3(num_classes=num_classes, backbone='DarkNet53')
model.train(
    num_epochs=270,
    train_dataset=train_dataset,
    train_batch_size=8,
    eval_dataset=eval_dataset,
    learning_rate=0.000125,
    lr_decay_epochs=[216, 243],
    warmup_steps=1000,
    warmup_start_lr=0.0,
    save_interval_epochs=5,
    save_dir='output/yolov3_darknet53',
    use_vdl=True)


# ## 4 模型预测
# 
# 使用模型进行预测，同时使用`pdx.det.visualize`将结果可视化，可视化结果将保存到`./output/yolov3_darknet53`下，其中`threshold`代表Box的置信度阈值，将Box置信度低于该阈值的框过滤不进行可视化。
# > 此处请在AIStudio Notebook页面的右上角菜单，选择重启执行器，以释放显存，重新加载训练好的模型

# In[ ]:


import paddlex as pdx
model = pdx.load_model('output/yolov3_darknet53/best_model')
image_name = 'insect_det/JPEGImages/0217.jpg'
result = model.predict(image_name)
pdx.det.visualize(image_name, result, threshold=0.5, save_dir='./output/yolov3_darknet53')


# 可视化结果如下所示：  
# <img src="https://ai-studio-static-online.cdn.bcebos.com/6d6c25bf26bd4390ab4f92e79cf3bca1bc0e8a975cc7475b9386a5cd1194aced" width="50%" height="50%">
# 
# 
# ## 5. 反馈
# PaddleX用户QQ群: 1045148026 (手机QQ扫描如下二维码快速加入)  
# <img src="https://ai-studio-static-online.cdn.bcebos.com/fef04deabad440e4918c891e23fce8d51b64c0c466dc4f0980ea42c3a0d6549a" width="15%" height="15%">
